import os  # 导入用于操作系统功能的模块
import torch   # 导入PyTorch深度学习库
import torch.nn as nn  # 导入神经网络模块
import torch.nn.functional as F   # 导入神经网络函数库
import matplotlib.pyplot as plt  # 导入用于绘图的matplotlib库
from torchsummary import summary

# Dice损失函数的定义
class DiceLoss(nn.Module):
    def __init__(self):
        # 用于计算Dice损失的PyTorch模块
        super(DiceLoss, self).__init__()

    # 定义模型的前向传播过程，接受张量pred、label、平滑系数smooth作为输入，返回损失值
    def forward(self, pred, label, smooth=1):
        pred = F.sigmoid(pred)  # 对预测结果进行Sigmoid变换，将其转换到0到1之间
        pred = pred.view(-1)  # 将预测结果展平为一维张量
        label = label.view(-1)  # 将真实标签展平为一维张量

        # 计算pred和label的交集
        intersection = (pred * label).sum()
        # 计算pred和label的并集
        union = pred.sum() + label.sum()
        # 根据前述的dice损失计算公式，加入平滑因子计算交并比
        dice = (2. * intersection + smooth) / (union + smooth)
        return 1 - dice

# 结合Dice损失和Binary Cross Entropy损失的DiceBCELoss损失函数
class DiceBCELoss(nn.Module):
    def __init__(self):
        super(DiceBCELoss, self).__init__()

    def forward(self,
                inputs: torch.Tensor,    # 模型的预测输出
                targets: torch.Tensor,   # 真实标签
                smooth: int = 1          # 平滑项
                ) -> torch.Tensor:       # 返回类型为Tensor
        # 对模型输出进行Sigmoid激活
        inputs = F.sigmoid(inputs)

        # 将预测值和真实标签转换为一维向量
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        # 计算交集
        intersection = (inputs * targets).sum()
        # 计算 Dice 损失      
        dice_loss = 1 - (2. * intersection + smooth) / (inputs.sum() + targets.sum() + smooth)
        # 计算 Binary Cross Entropy 损失
        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
        # 将 Dice 损失和 BCE 损失相加得到最终损失
        Dice_BCE = BCE + dice_loss
        return Dice_BCE


# 处理类别不平衡问题的损失函数
class FocalLoss(nn.Module):
    def __init__(self):
        super(FocalLoss, self).__init__()

    def forward(self, pred, label, alpha=0.8, gamma=2):
        pred = F.sigmoid(pred)  # 对预测值进行Sigmoid激活
        # 将pred和label展开为一维
        pred = pred.view(-1)
        label = label.view(-1)

        bce = F.binary_cross_entropy(pred, label, reduction='mean')  # 计算BCE损失
        bce = torch.exp(-bce)  # 对BCE进行指数转换

        focal_loss = alpha * (1 - bce) ** gamma * bce  # 计算Focal Loss
        return focal_loss  # 返回Focal Loss损失值

# 将损失值记录到文件中
class LossWriter():
    def __init__(self, save_dir):
        self.save_dir = save_dir  # 指定保存目录

    # 将损失名称 loss_name、损失值 loss 和迭代次数 i 追加到一个以损失名称命名的 .txt 文件中
    def add(self, loss_name, loss, i):
        with open(os.path.join(self.save_dir, loss_name + ".txt"), mode="a") as f:
            term = str(i) + " " + str(loss) + "\n"   # 格式化损失值和迭代次数
            f.write(term)  # 将损失值和迭代次数写入文件
            f.close()

# 绘制损失函数的函数
def plot_loss(txt_name, x_label, y_label, title, legend, font_size, fig_size, save_name):
    """
    损失函数绘图代码
    """
    all_i = []  # 存储迭代次数
    all_val = []  # 存储损失值
    with open(txt_name, "r") as f:  # 打开损失值文件
        all_lines = f.readlines()  # 读取txt文件中的所有行
        for line in all_lines:  # 遍历每一行
            sp= line.split(" ")  # 按空格拆分
            i = int(sp[0])  # 获取迭代次数
            val = float(sp[1])  # 获取损失值
            all_i.append(i)  # 存储迭代次数
            all_val.append(val)  # 存储损失值
    # 绘图以及参数指定
    plt.figure(figsize=(6, 4))  # 设置图形大小
    plt.plot(all_i, all_val)   # 绘制损失函数图像
    plt.xlabel(x_label, fontsize=font_size)  # 设置X轴标签和字体大小
    plt.ylabel(y_label, fontsize=font_size)  # 设置Y轴标签和字体大小
    if legend:
        plt.legend(legend, fontsize=font_size)  # 设置图例和字体大小
    plt.title(title, fontsize=font_size)  # 设置图像标题和字体大小
    plt.tick_params(labelsize=font_size)  # 设置刻度字体大小
    plt.savefig(save_name, dpi=200, bbox_inches = "tight")  # 保存图像
    plt.show()


if __name__ == "__main__":
    # 调用 plot_loss 函数加载损失数据并生成关于 BCE 损失的图像
    plot_loss(txt_name="results/loss/bce_loss.txt", x_label="iteration",
              y_label="loss value", title="Loss of BCE on UNet",
              legend=None, font_size=15, fig_size=(10, 10),
              save_name="unet_BCE_loss.png")